# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: beomsoo

import datetime
import re
import json
import os
import itertools
import zipfile
import fire
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from coin.base.query_util import (query_klines, query_pta)
from coin.support.market_quotes.logic.query_util import \
  _get_api_version_for_request
from reportlab.pdfgen.canvas import Canvas
from reportlab.platypus import Table
from reportlab.lib.utils import ImageReader

crit = {"market_share": 1.0,
    "one_tick_bp_fcs": 5.0,
    "one_tick_bp_ref": 10.0}
marker = itertools.cycle(('o','v','^','s','D'))
color = itertools.cycle(
    ('#1f77b4','#ff7f0e','#2ca02c',
    '#d62728','#9467bd','#8c564b',
    '#e377c2','#7f7f7f','#bcbd22'))
zscore_threshold = 4.0

def get_tick_size(market, exchange, symbol, avg_price):
  pi_dir = 'data/coin2/product_info'
  api = _get_api_version_for_request(market, exchange, symbol)
  mea = f'{market}.{exchange}.{api}'
  json_dir = os.path.join(pi_dir, mea+'.json')

  with open(json_dir, 'r') as f:
    data = json.load(f)
    tick_size = None
    for pi in data['product_infos']:
      if pi['symbol'] == symbol:
        for pts in pi['price_ticksize']:
          if exchange == 'Upbit' or exchange == 'Bithumb':
            price_multiplier = 1300
          elif exchange == 'Bitbank':
            price_multiplier = 140
          else:
            price_multiplier = 1
          if pts['from'] <= avg_price * price_multiplier and avg_price * price_multiplier < pts['to']:
            tick_size = pts['value'] / price_multiplier
        break
    return tick_size if tick_size != None else np.nan

def get_pta_df(start_date, period, strat_regex='.*lm.*'):
  start_date_dt = datetime.datetime.strptime(str(start_date), '%Y%m%d').date()
  end_date_dt = start_date_dt + datetime.timedelta(days=(period-1))
  end_date = end_date_dt.strftime('%Y%m%d')
  cutoff_days = 5
  if period > cutoff_days:
    curr_date_dt = start_date_dt
    pta_df = pd.DataFrame()
    while curr_date_dt < end_date_dt:
      next_date_dt = curr_date_dt + datetime.timedelta(days=(cutoff_days-1))
      if next_date_dt > end_date_dt:
        next_date_dt = end_date_dt
      df = query_pta(curr_date_dt, next_date_dt, business_units=['Coin'], pta_type='SYMBOL')
      pta_df = pd.concat([pta_df, df])
      curr_date_dt = curr_date_dt + datetime.timedelta(days=cutoff_days)
  else:
    pta_df = query_pta(start_date_dt, end_date_dt, business_units=['Coin'], pta_type='SYMBOL')
  def re_filter(x):
    return bool(re.search(strat_regex,x['strategy_name']))
  pta_df = pta_df[pta_df.apply(re_filter, axis=1)]
  pta_df['pnl'] = np.where(
    pta_df['strategy_name'].str.contains('spot'), pta_df['pnl_from_fill_in_usd'],
    np.where(pta_df['strategy_name'].str.contains('swap'), pta_df['pnl_from_balance_in_usd'], np.nan)
  )
  dft_group = ['market_type', 'exchange', 'symbol', 'strategy_name']
  valid_col = ['pnl', 'fill_pq_in_usd', 'fill_pq', 'fill_qty']
  pta_df[valid_col] = pta_df[valid_col].astype('float')
  td_df = pta_df[list(set(dft_group + ['trading_date']))]
  td_df = td_df.groupby(dft_group)
  date_df = td_df.min().reset_index().rename(columns={'trading_date': 'start_date'})
  date_df = date_df.merge(td_df.max().reset_index().rename(columns={'trading_date': 'end_date'}))
  date_df = date_df.merge(td_df.count().reset_index().rename(columns={'trading_date': 'num_dates'}))
  pta_df = pta_df[list(set(dft_group + valid_col))]
  sum_df = pta_df.groupby(dft_group).sum().reset_index()
  sum_df = pd.merge(sum_df, date_df, how='left')
  def get_avg_price(x):
    if 'coin' in x['strategy_name']:
      return x['fill_pq_in_usd'] / x['fill_pq'] if x['fill_pq'] != 0 else np.nan
    else:
      return x['fill_pq'] / x['fill_qty'] if x['fill_qty'] != 0 else np.nan
  sum_df['avg_price'] = sum_df.apply(get_avg_price, axis=1)
  sum_df['ret_bp'] = sum_df['pnl'] / sum_df['fill_pq_in_usd'] * 1e+4
  sum_df.dropna(subset=sum_df.columns, inplace=True)

  klines = query_klines(start_date_dt, end_date_dt)
  klines['external_symbol'] = klines['kline_dict'].apply(lambda x: json.loads(x).get('symbol'))
  klines['market_turnover_in_usd'] = klines['kline_dict'].apply(
      lambda x: (json.loads(x).get('klines', [{}]))[0].get('turnover_in_usd'))
  rklines = klines[['market_type','exchange','external_symbol','market_turnover_in_usd']]
  rklines = rklines.groupby(['market_type', 'exchange', 'external_symbol']).sum().reset_index()
  rklines.rename(columns={'external_symbol': 'symbol'}, inplace=True)
  res_df = pd.merge(sum_df, rklines, how='left')
  res_df['market_share'] = res_df['fill_pq_in_usd'] * 100.0 / res_df['market_turnover_in_usd']

  # get tick size of each symbol
  config_dir = 'python/coin_deploy/lm_agg2'
  grp_df = res_df.groupby('strategy_name')
  for strat_name, df in grp_df:
    for idx, symb in zip(df.index, df['symbol']):
      # fcs
      res_df.loc[idx, 'tick_size_fcs'] = get_tick_size(
          res_df.loc[idx, 'market_type'],
          res_df.loc[idx, 'exchange'],
          symb,
          res_df.loc[idx, 'avg_price']
      )
      res_df.loc[idx, 'one_tick_bp_fcs'] = \
          res_df.loc[idx, 'tick_size_fcs'] / res_df.loc[idx, 'avg_price'] * 1e4

      '''
      TODO(beomsoo): handle multiref
      # ref
      def get_mep(ref):
        if ref == 'bnref':
          return ('Spot', 'Binance', '-USDT')
        elif ref == 'bnsref':
          return ('Futures', 'Binance', '-USDT.PERPETUAL')
        elif ref == 'ubref':
          return ('Spot', 'Upbit', '-KRW')
        elif ref == 'cbref':
          return ('Spot', 'Gdax', '-USD')
        elif ref == 'okqref':
          return ('Futures', 'Okex', '-USD.QUARTER')
        elif ref == 'bnbusdref':
          return ('Spot', 'Binance', '-BUSD')
        elif ref == 'ftspotref':
          return ('Spot', 'Ftx', '-USD')
      for (root, dirs, files) in os.walk(config_dir):
        if len(files) > 0:
          for file_name in files:
            if strat_name + '.json' == file_name:
              with open(os.path.join(root,file_name), 'r') as f_conf:
                data = json.load(f_conf)
                prof_name = data['products'][str(list(data['products'])[0])]['profile_name']
                ref = prof_name.split('_')[2]
                (market, exchange, product_name) = get_mep(ref)
                res_df.loc[idx, 'tick_size_ref'] = get_tick_size(
                    market,
                    exchange,
                    symb.split('-')[0]+product_name,
                    res_df.loc[idx, 'avg_price']
                )
                if np.isnan(res_df.loc[idx, 'tick_size_ref']):
                  if symb.split('-')[0] == 'SHIB':
                    res_df.loc[idx, 'tick_size_ref'] = get_tick_size(
                        market,
                        exchange,
                        '1000SHIB'+product_name,
                        res_df.loc[idx, 'avg_price'] * 1000
                    ) / 1000
                  elif symb.split('-')[0] == '1000SHIB':
                    res_df.loc[idx, 'tick_size_ref'] = get_tick_size(
                        market,
                        exchange,
                        'SHIB'+product_name,
                        res_df.loc[idx, 'avg_price'] / 1000
                    ) * 1000
                res_df.loc[idx, 'one_tick_bp_ref'] = \
                    res_df.loc[idx, 'tick_size_ref'] / res_df.loc[idx, 'avg_price'] * 1e4
      '''

  return res_df

def plot_pta(start_date, period, directory,
             strat_regex='.*lm.*', x_name=None, y_name='ret_bp', export_csv=False):
  start_date_dt = datetime.datetime.strptime(str(start_date), '%Y%m%d').date()
  end_date_dt = start_date_dt + datetime.timedelta(days=(period-1))
  end_date = end_date_dt.strftime('%Y%m%d')
  assert end_date_dt < datetime.date.today(), f"end_date must be no later than yesterday"
  print(f"{start_date} - {end_date}")
  print(f"{strat_regex}")
  print(f"{x_name} VS {y_name}")
  plot_dir = os.path.join(directory, str(start_date) + '-' + end_date)
  if not os.path.exists(plot_dir):
    os.makedirs(plot_dir)
  df = get_pta_df(start_date, period, strat_regex)
  if export_csv:
    df.sort_values(by=[y_name], ascending=False).to_csv(os.path.join(plot_dir, "data.csv"), index=False)
  grp_df = df.groupby('strategy_name')

  # total
  if not os.path.exists(os.path.join(plot_dir, x_name)):
    os.makedirs(os.path.join(plot_dir, x_name))
  plt.figure(figsize=(10,5))
  min_date = None; max_date = None
  # If period > 3 days, only symbols that have traded more than 30% during the period are plotted
  for strat_name, df in grp_df:
    each_min_date = df['start_date'].min() if period <= 3 or (period > 3 and (df['num_dates'] >= 0.3 * period).sum() == 0) \
        else df[df['num_dates'] >= 0.3 * period]['start_date'].min()
    min_date = each_min_date if min_date == None else min(min_date, each_min_date)
    each_max_date = df['end_date'].max() if period <= 3 or (period > 3 and (df['num_dates'] >= 0.3 * period).sum() == 0) \
        else df[df['num_dates'] >= 0.3 * period]['end_date'].max()
    max_date = each_max_date if max_date == None else max(max_date, each_max_date)
    clr = next(color); mrkr = next(marker)
    x = df[x_name] if period <= 3 else df[df['num_dates'] >= 0.3 * period][x_name]
    y = df[y_name] if period <= 3 else df[df['num_dates'] >= 0.3 * period][y_name]
    plt.plot(x, y, color=clr, marker=mrkr, linestyle='', label=strat_name)
  if min_date != None and max_date != None:
    plt.axvline(x=crit[x_name], ymin=-100, ymax=100, color='r', linestyle='-', linewidth=2)
    plt.title(min_date.strftime("%Y%m%d")+'-'+max_date.strftime("%Y%m%d"))
    plt.xlabel(x_name)
    plt.ylabel(y_name)
    plt.legend(bbox_to_anchor=(1.04,1.06), loc="upper left", fontsize=6.5)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(os.path.join(plot_dir,x_name+'/total.png'))
    plt.close()
  # each strat
  if not os.path.exists(os.path.join(plot_dir, x_name+'/each_strat')):
    os.makedirs(os.path.join(plot_dir, x_name+'/each_strat'))
  for strat_name, df in grp_df:
    plt.figure(figsize=(10,5))
    min_date = None; max_date = None
    for symb in df['symbol']:
      if period <= 3 or (period > 3 and df.loc[df['symbol'] == symb, 'num_dates'].iloc[0] >= 0.3 * period):
        min_date = df.loc[df['symbol'] == symb, 'start_date'].iloc[0] if min_date == None \
            else min(min_date, df.loc[df['symbol'] == symb, 'start_date'].iloc[0])
        max_date = df.loc[df['symbol'] == symb, 'end_date'].iloc[0] if max_date == None \
            else max(max_date, df.loc[df['symbol'] == symb, 'end_date'].iloc[0])
        plt.plot(df.loc[df['symbol'] == symb, x_name],
            df.loc[df['symbol'] == symb, y_name],
            color=next(color), marker=next(marker), linestyle='', label=symb)
    if min_date != None and max_date != None:
      plt.axvline(x=crit[x_name], ymin=-100, ymax=100, color='r', linestyle='-', linewidth=2)
      plt.title(min_date.strftime("%Y%m%d")+'-'+max_date.strftime("%Y%m%d"))
      plt.xlabel(x_name)
      plt.ylabel(y_name)
      plt.legend(bbox_to_anchor=(1.04,1.06), loc="upper left", fontsize=6.5)
      plt.grid(True)
      plt.tight_layout()
      plt.savefig(os.path.join(plot_dir,x_name+'/each_strat/'+strat_name+'.png'))
      plt.close()

def performance_eval(start_date=None, period=1, strat_regex='.*lm.*', max_num=7, export_pdf=False, directory='performance_eval/'):
  assert max_num >= 1, f"max_num should be a positive number (max_num={max_num})"
  if start_date == None: # set recent 'period' dates
    end_date_dt = datetime.datetime.now() - datetime.timedelta(days=1)
    start_date_dt = end_date_dt - datetime.timedelta(days=(period-1))
    start_date = start_date_dt.strftime('%Y%m%d')
    end_date = end_date_dt.strftime('%Y%m%d')
  else:
    start_date_dt = datetime.datetime.strptime(str(start_date), '%Y%m%d').date()
    end_date_dt = start_date_dt + datetime.timedelta(days=(period-1))
    end_date = end_date_dt.strftime('%Y%m%d')
    assert end_date_dt < datetime.date.today(), f"end_date must be no later than yesterday"

  print(f"{start_date} - {end_date}")
  df = get_pta_df(start_date, period, strat_regex)
  grp_df = df.groupby('strategy_name')
  sum_df = grp_df.sum()[['pnl', 'fill_pq_in_usd', 'market_turnover_in_usd']]
  sum_df['ret_bp'] = sum_df['pnl'] / sum_df['fill_pq_in_usd'] * 1e4
  sum_df['market_share'] = sum_df['fill_pq_in_usd'] / sum_df['market_turnover_in_usd'] * 1e2
  strat_summary = sum_df.sort_values(by='ret_bp', ascending=False)[['pnl', 'ret_bp', 'market_share']]
  strat_summary.rename(columns = {'pnl': 'pnl (USD)',
                                  'ret_bp': 'net_ret (bp)',
                                  'market_share': 'market share (%)'}, inplace=True)
  df.insert(len(df.columns), 'score', ((df['pnl'] - df['pnl'].mean())
       /df['pnl'].std(ddof=0)+(df['ret_bp']-df['ret_bp'].mean())/df['ret_bp'].std(ddof=0)))
  top_summary = df[df['num_dates'].astype('int32')>=0.5*period].sort_values(by='score', ascending=False).reset_index(drop=True).head(max_num) \
      [['strategy_name', 'symbol', 'pnl', 'ret_bp', 'market_share', 'num_dates']]
  top_summary.rename(columns = {'strategy_name': 'strategy',
                                'pnl': 'pnl (USD)',
                                'ret_bp': 'net_ret (bp)',
                                'market_share': 'market share (%)',
                                'num_dates': 'trading days'}, inplace=True)
  bottom_summary = df[df['num_dates'].astype('int32')>=0.5*period].sort_values(by='score', ascending=True).reset_index(drop=True).head(max_num) \
      [['strategy_name', 'symbol', 'pnl', 'ret_bp', 'market_share', 'num_dates']]
  bottom_summary.rename(columns = {'strategy_name': 'strategy',
                                   'pnl': 'pnl (USD)',
                                   'ret_bp': 'net_ret (bp)',
                                   'market_share': 'market share (%)',
                                   'num_dates': 'trading days'}, inplace=True)
  prob_pnl = pd.DataFrame()
  for strat_name, df_each in grp_df:
    if len(df_each) < 5:
      continue
    for col in ['market_share', 'ret_bp', 'pnl']:
      col_zscore = col + '_zscore'
      df_each.insert(len(df_each.columns), col_zscore, (df_each[col] - df_each[col].mean()) / df_each[col].std(ddof=0))
    prob_pnl = prob_pnl.append(df_each[df_each['pnl_zscore']<-zscore_threshold]
        [['strategy_name', 'symbol', 'pnl', 'pnl_zscore', 'ret_bp', 'market_share']])
  prob_pnl = prob_pnl.sort_values(by='pnl_zscore', ascending=True).reset_index(drop=True)
  prob_pnl.rename(columns = {'strategy_name': 'strategy',
                             'pnl': 'pnl (USD)',
                             'pnl_zscore': 'pnl z-score',
                             'ret_bp': 'net_ret (bp)',
                             'market_share': 'market share (%)',
                                   'num_dates': 'trading days'}, inplace=True)
  print(f"1. Strategy Summary (sorted by net_ret)")
  print(strat_summary)
  print(f"\n2. Top {max_num} Symbols")
  print(top_summary)
  print(f"\n3. Bottom {max_num} Symbols")
  print(bottom_summary)
  print(f"\n4. Problematic Symbols (z-score threshold=-{zscore_threshold})")
  print(prob_pnl)
  if export_pdf:
    pdf_dir = directory
    if not os.path.exists(pdf_dir):
      os.makedirs(pdf_dir)
    # 1. Strategy Summary (sorted by net_ret)
    canvas = Canvas(pdf_dir+f"{strat_regex}_{start_date}-{end_date}"+'.pdf')
    canvas.setPageSize((600, 1300))
    canvas.setFont('Helvetica-Bold', 20)
    canvas.drawString(40, 1240, "1. Strategy Summary (sorted by net_ret)")
    list_summary = [[''] + strat_summary.columns.tolist()]
    for idx, row in strat_summary.iterrows():
      list_summary.append([idx, f"{row[0]:.2f}", f"{row[1]:.2f}", f"{row[2]:.3f}"])
    table = Table(list_summary)
    table.wrapOn(canvas, 200, 200)
    table.drawOn(canvas, 40, 60)
    canvas.showPage()
    # 2. Top {max_num} Strats
    strat_summary = sum_df.sort_values(by='pnl', ascending=False)
    canvas.setPageSize((800, 100))
    canvas.setFont('Helvetica-Bold', 20)
    canvas.drawString(40, 40, "2. TOP " + f"{max_num}" + " strategies: in order of pnl")
    canvas.showPage()
    if period == 1:
      live_period = 1
    elif period > 1 and period <= 5:
      live_period = 5
    elif period > 5 and period <= 10:
      live_period = 10
    else:
      live_period = 20
    live_start_date_dt = end_date_dt - datetime.timedelta(days=(live_period-1))
    live_end_date_dt = end_date_dt
    while 1:
      live_start_date = live_start_date_dt.strftime('%Y%m%d')
      live_end_date = live_end_date_dt.strftime('%Y%m%d')
      pplot_dir = os.path.join(
          "/remote/iosg/strat-1/buckets/analysis.derived.coin/live/plot/pplot/",
          live_start_date+"-"+live_end_date,
          "merged")
      if os.path.isdir(pplot_dir):
        break
      live_start_date_dt = live_start_date_dt - datetime.timedelta(days=1)
      live_end_date_dt = live_end_date_dt - datetime.timedelta(days=1)
    canvas.setPageSize((800, 750))
    cnt = 1
    for strat_name, row in strat_summary.iterrows():
      for file_name in os.listdir(pplot_dir):
        if strat_name in file_name:
          canvas.setFont('Helvetica', 10)
          canvas.drawString(10, 740, str(cnt)+') http://files.ap.ngrok.io'+pplot_dir+file_name)
          im = ImageReader(os.path.join(pplot_dir,file_name))
          canvas.drawImage(im, 100, 10, width = 644, height = 730, mask='auto')
          canvas.showPage()
          break
      if cnt == max_num:
        break
      cnt = cnt + 1
    # 3. Bottom {max_num} Strats
    strat_summary = sum_df.sort_values(by='pnl', ascending=True)
    canvas.setPageSize((800, 100))
    canvas.setFont('Helvetica-Bold', 20)
    canvas.drawString(40, 40, "3. BOTTOM " + f"{max_num}" + " strategies: in order of pnl")
    canvas.showPage()
    canvas.setPageSize((800, 750))
    cnt = 1
    for strat_name, row in strat_summary.iterrows():
      for file_name in os.listdir(pplot_dir):
        if strat_name in file_name:
          canvas.setFont('Helvetica', 10)
          canvas.drawString(10, 740, str(cnt)+') http://files.ap.ngrok.io'+pplot_dir+file_name)
          im = ImageReader(os.path.join(pplot_dir,file_name))
          canvas.drawImage(im, 100, 10, width = 644, height = 730, mask='auto')
          canvas.showPage()
          break
      if cnt == max_num:
        break
      cnt = cnt + 1
    # 4. Problematic symbols
    list_summary = [prob_pnl.columns.tolist()]
    for idx, row in prob_pnl.iterrows():
      list_summary.append([row[0], row[1], f"{row[2]:.2f}", f"{row[3]:.3f}", f"{row[4]:.3f}", f"{row[5]:.3f}"])
    canvas.setPageSize((800, 600))
    canvas.setFont('Helvetica-Bold', 20)
    canvas.drawString(40, 560, "4. Problematic symbols: symbols with pnl z-score < " + f"-{zscore_threshold}" + " within a strategy")
    table = Table(list_summary)
    table.wrapOn(canvas, 200, 200)
    table.drawOn(canvas, 40, 40)
    canvas.showPage()
    canvas.setPageSize((800, 750))
    pplot_dir = pplot_dir.replace("each", "merged")
    cnt = 1
    for idx, row in prob_pnl.iterrows():
      for file_name in os.listdir(pplot_dir):
        if row['strategy']+'_' in file_name:
          canvas.setFont('Helvetica', 10)
          canvas.drawString(10, 740, str(cnt)+') http://files.ap.ngrok.io'+pplot_dir+file_name)
          im = ImageReader(os.path.join(pplot_dir,file_name))
          canvas.setFont('Helvetica', 14)
          canvas.drawString(10, 720, f"{cnt}. symbol: {row[1]}, z-score: {row[3]:.3f}")
          canvas.drawImage(im, 100, -12, width = 644, height = 730, mask='auto')
          canvas.showPage()
          break
      cnt = cnt + 1
    canvas.save()
    print('PDF EXPORTED: http://files.ap.ngrok.io'+os.getcwd()+'/'+directory+f"{strat_regex}_{start_date}-{end_date}"+'.pdf')

if __name__ == '__main__':
  fire.Fire()

'''
Usage
./pyrunner python/experimental/bumsoo/strat_pta.py plot_pta --start_date=20220620 --period=10 --directory=python/experimental/bumsoo/ --strat_regex=.*lm_focus_st.* --x_name=market_share --y_name=ret_bp --export_csv
    x_name, y_name: ['pnl', 'fill_pq_in_usd', 'fill_pq', 'fill_qty',
                 'market_turnover_in_usd', 'ret_bp', 'market_share', 'one_tick_bp_fcs', 'one_tick_bp_ref']

./pyrunner python/experimental/bumsoo/strat_pta.py performance_eval --start_date=20220601 --period=20 --max_num=7 --strat_regex=".*lm_focus_(st|volume).*" --export_pdf --directory=python/experimental/bumsoo/pnl_plot/
'''
